
import cv2
import numpy as np
from scipy.ndimage import filters
from images.rof import ImageBaseFunc
from images.base import AVBase
from images.base import FaceDetection
import sys
import os

def do_test_facedetection(fd, im):
    return fd.detect(im)

def do_test_cornerharris(im):
    gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY)
    gray = np.float32(gray)
    dst = cv2.cornerHarris(gray, 2, 23, 0.04)
    im[dst > 0.01 * dst.max()] = [0, 0, 255]
    return im

def do_test_harris(im):
    harrisim = ImageBaseFunc.jk_compute_harris_response(im)
    filtered_coords = ImageBaseFunc.jk_get_harris_points(harrisim, 6)
    for p in filtered_coords:
        cv2.cvCircle(im, p, 3, cv2.CV_RGB(255, 0, 0), 1, 8, 3)
    return im

def do_test_filter(im):
    img = cv2.cvtColor(im, cv2.COLOR_RGB2GRAY)
    return filters.gaussian_filter(img, 10)

def do_test_edge(b, im):
    dst = im.copy()
    b.stoke_edges(im, dst)
    return dst

def do_test_canny(im):
    return cv2.Canny(im, 50, 300)

count = 0

def do_with_base(b, im):
    # return do_test_harris(im)
    global count
    count += 1
    if count % 25 < 20:
        return do_test_canny(im)
    else:
        return im
    # return do_test_cornerharris(im)
    # return b.line_detect_image(im)s
    # return b.contour_image(im)
    # return do_test_edge(b, im)
    # return b.do_sift(im)

def do_with_facedetection(f, im):
    return do_test_facedetection(f, im)

def start_video():
    cam = cv2.VideoCapture(0)
    cv2.namedWindow("fd")
    avbase = AVBase()
    # fd = FaceDetection()
    while cam.isOpened():
        flag, frame = cam.read()
        if flag == True:
            # newimg = do_with_facedetection(fd, frame)
            newimg = do_with_base(avbase, frame)
            cv2.imshow("fd", newimg)
            if (cv2.waitKey(10) == 27):
                cam.release()
                break
    cv2.destroyWindow("fd")

def start_video_file(file):
    cam = cv2.VideoCapture(file)
    avbase = AVBase()
    fd = FaceDetection()
    frame_index = 0
    img_index = 0
    while cam.isOpened():
        flag, frame = cam.read()
        if flag == True:
            rects = fd.detect_rects(frame)
            if len(rects) > 0:
                frame_index = frame_index + 1
                img_index = 0
                print("Find face counts ", len(rects))
                for rct in rects:
                    img_index = img_index + 1
                    img = frame[rct[0]:rct[0] + rct[2], rct[1]:rct[1] + rct[3]]
                    bname = os.path.basename(file).split(".")[0]
                    saveto = os.path.dirname(file) + "\\out\\" + bname + "-fd-" + str(frame_index) + "-" + str(img_index) + ".jpg"
                    print("saveto: ", saveto)
                    cv2.imwrite(saveto, img)
        else:
            break

if __name__ == "__main__":
    if len(sys.argv) > 1:
        print("do with file ", sys.argv[0])
        start_video_file(sys.argv[1])
    else:
        start_video()